import os
import re
import paramiko
import logging

 
class NpuMonitor:
    def __init__(self, hostname, username, password):
        self.hostname = hostname
        self.username = username
        self.password = password

    def get_structured_npu_info(hostname, username, password):
        try:
            # 连接到服务器
            client = paramiko.SSHClient()
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            client.connect(hostname, username=username, password=password)

            # 执行命令
            stdin, stdout, stderr = client.exec_command('npu-smi info')
            output = stdout.read().decode()
            error = stderr.read().decode()
            
            if error:
                print(f"错误: {error}")
                return []

            # 解析输出
            lines = output.strip().split('\n')[2:]
            npu_list = []
            for line in lines:
                match = re.match(r'\s*(\d+)\s+\w+\s+(\d+)\s+([\d.]+)\s+(\d+)\s+(\d+)/(\d+)\s+(\d+)\s+(\d+)\s+(\d+)', line)
                if match:
                    npu_list.append({
                        "NPU ID": match.group(1),
                        "Chip ID": match.group(2),
                        "pwr(W)": match.group(3),
                        "Temp(C)": match.group(4),
                        "AI Core(%)": match.group(5),
                        "AI Cpu(%)": match.group(6),
                        "Ctrl Cpu(%)": match.group(7),
                        "Memory(%)": match.group(8),
                        "Memory BW(%)": match.group(9)
                    })
            return lines
              
        except Exception as e:
            print(f"操作失败: {str(e)}")
            return []
        finally:
            client.close()


if __name__ == "__main__":
    npu_info = NpuMonitor.get_structured_npu_info("177.10.8.2", "root", "weique888")
    print("NPU 状态:", npu_info)